//! task 相关操作

use alloc::{format, string::String, sync::Arc, vec::Vec};
use core::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicUsize, Ordering};

use bsp_define::smp::NR_CPUS;
use task_define::TaskBase;

use crate::{
    processor::{nr_cpus, this_processor_id},
    sched::task::{
        pid::{Pid, pid_init},
        stack::THREAD_STACK_SIZE,
        thread::{TASK_NEED_RESCHED, ThreadStruct},
    },
    space::{
        addr::{Vaddr, init_stack_base},
        kalloc::{GfpFlags, kfree, kmalloc},
        mm::MmStruct,
    },
    time::timekeeping::ktime_get,
};

mod pid;
pub mod stack;
pub(crate) mod thread;

/// 任务结构体
#[repr(C)]
pub struct Task {
    // TaskBase必须在第一个位置, 汇编需要使用相关偏移
    task_base: TaskBase,
    thread: ThreadStruct,
    badge: u64,
    name: String,
    cpu: AtomicUsize,
    on_cpu: AtomicBool,
    mm: Option<Arc<MmStruct>>,
    pid: Pid,
    state: AtomicU32,
    sched_class: AtomicU32,
    n_nr_switch: AtomicU64,
    nv_nr_switch: AtomicU64,
}

/// 调度类型
pub type TaskSchedClass = u32;

/// `RT`实时任务
pub const TASK_SCHED_CLASS_RT: TaskSchedClass = 0x0000_0001;
/// `CFS`完全公平任务
pub const TASK_SCHED_CLASS_CFS: TaskSchedClass = 0x0000_0002;
/// `IDLE`内核空闲任务
pub const TASK_SCHED_CLASS_IDLE: TaskSchedClass = 0x0000_0004;
/// 未设置调度类型任务
pub const TASK_SCHED_CLASS_UNKNOWN: TaskSchedClass = 0x1000_0000;

/// 任务状态
pub type TaskState = u32;

/// 新任务
pub const TASK_NEW: TaskState = 0x1000_0000;
/// 任务正在运行
pub const TASK_RUNNING: TaskState = 0x0000_0000;

impl PartialEq for Task {
    fn eq(&self, other: &Self) -> bool {
        self.badge == other.badge
    }
}

impl Task {
    const fn new() -> Self {
        Task {
            task_base: TaskBase::new(),
            thread: ThreadStruct::new(),
            badge: 0,
            name: String::new(),
            cpu: AtomicUsize::new(0),
            on_cpu: AtomicBool::new(false),
            mm: None,
            pid: Pid::new(),
            state: AtomicU32::new(TASK_NEW),
            sched_class: AtomicU32::new(TASK_SCHED_CLASS_UNKNOWN),
            n_nr_switch: AtomicU64::new(0),
            nv_nr_switch: AtomicU64::new(0),
        }
    }

    fn init(name: &str, cpu: usize, mm: Option<Arc<MmStruct>>, stack: usize) -> Self {
        let tsk = Self {
            task_base: TaskBase::new(),
            thread: ThreadStruct::new(),
            badge: ktime_get(),
            name: String::from(name),
            cpu: AtomicUsize::new(cpu),
            on_cpu: AtomicBool::new(false),
            mm,
            pid: Pid::create(),
            state: AtomicU32::new(TASK_NEW),
            sched_class: AtomicU32::new(TASK_SCHED_CLASS_UNKNOWN),
            n_nr_switch: AtomicU64::new(0),
            nv_nr_switch: AtomicU64::new(0),
        };
        tsk.task_base.set_stack(stack);
        tsk
    }
    /// 创建一个任务
    ///
    /// # Panics
    /// 内存不足将会 panic.
    pub fn create(name: &str, cpu: usize, mm: Option<Arc<MmStruct>>) -> Self {
        let stack_base = kmalloc(THREAD_STACK_SIZE, GfpFlags::Clean).unwrap().to_value();
        Task::init(name, cpu, mm, stack_base)
    }

    /// 任务名
    pub fn name(&self) -> &str {
        &self.name
    }

    /// 任务所在处理器
    pub fn cpu(&self) -> usize {
        self.cpu.load(Ordering::Relaxed)
    }

    /// 任务地址空间
    pub fn mm(&self) -> Option<Arc<MmStruct>> {
        self.mm.as_ref().map(|m| m.clone())
    }

    /// 设置任务抢占计数
    #[inline(always)]
    pub fn set_preempt(&self, val: u32) {
        self.task_base.set_preempt(val);
    }

    /// 获取任务抢占计数
    #[inline(always)]
    pub fn preempt(&self) -> u32 {
        self.task_base.preempt()
    }

    /// 增加任务抢占计数
    #[inline(always)]
    pub fn preempt_add(&self, val: u32) {
        self.task_base.preempt_add(val);
    }

    /// 减少任务抢占计数
    #[inline(always)]
    pub fn preempt_sub(&self, val: u32) {
        self.task_base.preempt_sub(val);
    }

    /// 增加一次任务抢占计数
    #[inline(always)]
    pub fn preempt_inc(&self) {
        self.preempt_add(1);
    }

    /// 减少一次任务抢占计数
    #[inline(always)]
    pub fn preempt_dec(&self) {
        self.preempt_sub(1);
    }

    /// 任务`pid`
    #[inline(always)]
    pub fn pid(&self) -> usize {
        self.pid.to_value()
    }

    /// 设置任务状态
    #[inline(always)]
    pub fn set_state(&self, state: TaskState) {
        self.state.store(state, Ordering::Relaxed);
    }

    /// 获取任务状态
    #[inline(always)]
    pub fn state(&self) -> TaskState {
        self.state.load(Ordering::Relaxed)
    }

    /// 设置任务调度类型
    #[inline(always)]
    pub fn set_sched_class(&self, class: TaskSchedClass) {
        self.sched_class.store(class, Ordering::Relaxed);
    }

    /// 获取任务调度类型
    #[inline(always)]
    pub fn sched_class(&self) -> TaskSchedClass {
        self.sched_class.load(Ordering::Relaxed)
    }

    /// 清任务需要重新调度标志
    #[inline(always)]
    pub fn clear_need_resched(&self) {
        self.task_base.clear_flags_bit(TASK_NEED_RESCHED);
    }

    /// 任务需要重新调度
    #[inline(always)]
    pub fn need_resched(&self) -> bool {
        self.task_base.test_flags_bit(TASK_NEED_RESCHED)
    }

    /// 徽章
    #[inline(always)]
    pub fn badge(&self) -> u64 {
        self.badge
    }

    /// 增加自愿调度计数
    #[inline(always)]
    pub fn inc_n_switch_count(&self) {
        self.n_nr_switch.fetch_add(1, Ordering::Relaxed);
    }

    /// 增加非自愿调度计数
    #[inline(always)]
    pub fn inc_nv_switch_count(&self) {
        self.nv_nr_switch.fetch_add(1, Ordering::Relaxed);
    }

    /// 设置在处理器上运行
    #[inline(always)]
    pub fn set_on_cpu(&self) {
        self.on_cpu.store(true, Ordering::Relaxed);
    }

    /// 清除在处理器上运行
    pub fn clear_on_cpu(&self) {
        self.on_cpu.store(false, Ordering::Relaxed);
    }

    /// 获取是否在处理器上运行
    #[inline(always)]
    pub fn on_cpu(&self) -> bool {
        self.on_cpu.load(Ordering::Relaxed)
    }
}

impl Drop for Task {
    fn drop(&mut self) {
        let stack_base = self.task_base.stack();
        kfree(Vaddr::from(stack_base));
    }
}

/// 当前任务引用
///
/// # Panics
/// 系统故障将会 panic.
#[inline(always)]
pub fn current() -> Arc<Task> {
    let this_task = unsafe { PERCPU_ENTRY_TASK[this_processor_id()].as_ref().unwrap() };
    this_task.clone()
}

/// 当前任务快速访问引用
#[inline(always)]
pub fn current_fast() -> &'static Task {
    cfg_if::cfg_if! {
        if #[cfg(target_arch = "aarch64")] {
            unsafe {
                let task: usize;
                core::arch::asm!(
                    "mrs {0}, sp_el0",
                    out(reg) task
                );
                &*(task as *const Task)
            }
        } else {
            unreachable!()
        }
    }
}

// 仅唤醒从核时获取对应处理器任务信息使用
#[inline]
pub(crate) fn get_idle_task_info(cpu: usize) -> (usize, usize) {
    unsafe {
        let ts = PERCPU_ENTRY_TASK_RAW[cpu];
        let task = &*(ts as *const Task);
        let stack = task.task_base.stack();
        (ts, stack)
    }
}

#[inline(always)]
#[allow(unused_variables)]
fn set_current_fast(ts: usize) {
    cfg_if::cfg_if! {
        if #[cfg(target_arch = "aarch64")] {
            unsafe {
                core::arch::asm!(
                    "msr sp_el0, {0}",
                    in(reg) ts
                );
            }
        } else {
            unreachable!()
        }
    }
}

static DUMMY_TASK: Task = Task::new();

pub(crate) fn early_task_init() {
    set_current_fast(&raw const DUMMY_TASK as usize);
}

fn root_task_init() -> usize {
    let tsk = Arc::new(Task::init(
        "root_server",
        0,
        Some(Arc::new(MmStruct::create().unwrap())),
        init_stack_base().to_value(),
    ));
    tsk.set_state(TASK_RUNNING);
    tsk.set_on_cpu();
    let raw_tsk = tsk.as_ref() as *const Task as usize;
    switch_entry_task(0, tsk);
    raw_tsk
}

#[allow(static_mut_refs)]
pub(crate) fn task_init() {
    unsafe {
        for _ in 0..nr_cpus() {
            PERCPU_ENTRY_TASK.push(None);
        }
    }

    pid_init();

    let root_tsk = root_task_init();

    for idx in 1..nr_cpus() {
        let idle = idle_create(idx);
        idle.set_state(TASK_RUNNING);
        idle.set_on_cpu();
        switch_entry_task(idx, idle);
    }

    let preempt = current_fast().preempt();
    set_current_fast(root_tsk);
    assert_eq!(preempt, current_fast().preempt());

    // ok. 切换用户地址空间
    current().mm().unwrap().switch_mm();
}

// NOTE: 只是创建 idle 任务, 不要创建重复了
pub(crate) fn idle_create(cpu: usize) -> Arc<Task> {
    Arc::new(Task::create(&format!("idle/{cpu}"), cpu, None))
}

#[inline(always)]
pub(crate) fn switch_entry_task(cpu: usize, tsk: Arc<Task>) {
    let raw_tsk = tsk.as_ref() as *const Task as usize;
    unsafe {
        PERCPU_ENTRY_TASK[cpu] = Some(tsk);
        PERCPU_ENTRY_TASK_RAW[cpu] = raw_tsk;
    }
}

// 记录当前任务地址
static mut PERCPU_ENTRY_TASK: Vec<Option<Arc<Task>>> = Vec::new();

// 记录当前任务原始指针
//
// # NOTE
// 仅用于汇编使用, 并且与`PERCPU_ENTRY_TASK`同步更新
#[unsafe(no_mangle)]
static mut PERCPU_ENTRY_TASK_RAW: [usize; NR_CPUS] = [0; NR_CPUS];
